#include "UdpServer.h"
#include <cstdio>

namespace LiteNetLib
{
    void UdpServer::getAddress(struct sockaddr *addr, std::string &outAddress)
    {
        char strBuff[INET6_ADDRSTRLEN];
        if (addr->sa_family == AF_INET)
        {
            struct sockaddr_in *addr_in = (struct sockaddr_in *)addr;
            char ip_str[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &addr_in->sin_addr, ip_str, sizeof(ip_str));

            sprintf(strBuff, "%s:%d", ip_str, addr_in->sin_port);
            outAddress = strBuff;
        }
        else if (addr->sa_family == AF_INET6)
        {
            struct sockaddr_in6 *addr_in6 = (struct sockaddr_in6 *)addr;
            char ip_str[INET6_ADDRSTRLEN];
            inet_ntop(AF_INET6, &addr_in6->sin6_addr, ip_str, sizeof(ip_str));

            sprintf(strBuff, "[%s]:%d", ip_str, addr_in6->sin6_port);
            outAddress = strBuff;
        }
    }

    void UdpServer::getIpStr(struct sockaddr *addr, std::string &outIpStr)
    {
        if (addr->sa_family == AF_INET)
        {
            struct sockaddr_in *addr_in = (struct sockaddr_in *)addr;
            char ip_str[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &addr_in->sin_addr, ip_str, sizeof(ip_str));
            outIpStr = ip_str;
        }
        else if (addr->sa_family == AF_INET6)
        {
            struct sockaddr_in6 *addr_in6 = (struct sockaddr_in6 *)addr;
            char ip_str[INET6_ADDRSTRLEN];
            inet_ntop(AF_INET6, &addr_in6->sin6_addr, ip_str, sizeof(ip_str));
            outIpStr = ip_str;
        }
    }

    int UdpServer::getPort(struct sockaddr *addr)
    {
        if (addr->sa_family == AF_INET)
        {
            struct sockaddr_in *addr_in = (struct sockaddr_in *)addr;
            return addr_in->sin_port;
        }
        else if (addr->sa_family == AF_INET6)
        {
            struct sockaddr_in6 *addr_in6 = (struct sockaddr_in6 *)addr;
            return addr_in6->sin6_port;
        }
        return -1;
    }

    UdpServer::UdpServer()
    {
#ifdef _WIN32
        if (WSAStartup(MAKEWORD(2, 2), &m_wsaData) != 0)
        {
            throw std::runtime_error("WSAStartup failed");
        }
#endif
    }

    UdpServer::~UdpServer()
    {
#ifdef _WIN32
        WSACleanup();
#endif

        stop();
    }

    int UdpServer::create(int port, const char *host)
    {
        char strBuff[16];

        struct addrinfo hints;
        struct addrinfo *serverinfo;
        struct addrinfo *p;

        std::memset(&hints, 0, sizeof(hints));
        hints.ai_family = AF_UNSPEC; // IPv4 or IPv6
        hints.ai_socktype = SOCK_DGRAM;
        hints.ai_flags = AI_PASSIVE;

        sprintf(strBuff, "%d", port);

        int rv = getaddrinfo(host, strBuff, &hints, &serverinfo);
        if (rv != 0)
        {
            throw std::runtime_error("getaddrinfo failed: " + std::string(gai_strerror(rv)));
        }

        for (p = serverinfo; p != nullptr; p = p->ai_next)
        {
            m_sockfd = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
            if (m_sockfd == INVALID_SOCKET)
            {
                continue;
            }
#ifdef WIN32
            int enable = 1;
            if (setsockopt(m_sockfd, SOL_SOCKET, SO_REUSEADDR, (const char *)&enable, sizeof(enable)) == SOCKET_ERROR)
            {
                closesocket(m_sockfd);
                throw std::runtime_error("setsockopt SO_REUSEADDR failed");
            }
#endif

            if (bind(m_sockfd, p->ai_addr, p->ai_addrlen) == -1)
            {
#ifdef WIN32
                closesocket(m_sockfd);
#else
                shutdown(m_sockfd, SHUT_RDWR);
#endif
                continue;
            }
            getAddress(p->ai_addr, m_address);
            break;
        }

        if (p == nullptr)
        {
            throw std::runtime_error("Failed to bind socket");
        }

#if WIN32
        u_long mode = 1;
        if (ioctlsocket(m_sockfd, FIONBIO, &mode) == SOCKET_ERROR)
        {
            throw std::runtime_error("ioctlsocket failed");
        }

        m_iocpPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, nullptr, 0, 0);
        if (m_iocpPort == nullptr)
        {
            throw std::runtime_error("CreateIoCompletionPort failed");
        }
        CreateIoCompletionPort(reinterpret_cast<HANDLE>(m_sockfd), m_iocpPort, 0, 0);

        memset(&(m_ioData.overlapped), 0, sizeof(WSAOVERLAPPED));
        m_ioData.wsabuf.buf = m_ioData.buffer;
        m_ioData.wsabuf.len = sizeof(m_ioData.buffer);
        m_ioData.operation = 0;
        m_ioData.remoteAddrLen = sizeof(m_ioData.remoteAddr);
        DWORD flags = 0;
        DWORD bytesReceived;
        WSARecvFrom(m_sockfd, &(m_ioData.wsabuf), 1, &bytesReceived, &flags, (SOCKADDR *)&(m_ioData.remoteAddr), &(m_ioData.remoteAddrLen), &(m_ioData.overlapped), NULL);
#else
        int flags = fcntl(m_sockfd, F_GETFL, 0);
        fcntl(m_sockfd, F_SETFL, flags | O_NONBLOCK);
#endif

        freeaddrinfo(serverinfo);
        return m_sockfd;
    }

    void UdpServer::receiveData()
    {


#if WIN32

        WSAOVERLAPPED overlapped;
        DWORD bytesTransferred;
        ULONG_PTR completionKey;
        LPOVERLAPPED lpOverlapped;

        IO_DATA *lpIoData = NULL;

        while (m_isRunning)
        {
            BOOL result = GetQueuedCompletionStatus(m_iocpPort, &bytesTransferred, &completionKey, &lpOverlapped, INFINITE);
            if (result && bytesTransferred > 0)
            {
                lpIoData = (IO_DATA *)lpOverlapped;
                if (lpIoData->operation == 0)
                {
                    if (onMessage)
                    {
                        onMessage(m_ioData.buffer, bytesTransferred, (struct sockaddr *)&m_ioData.remoteAddr);
                    }
                }
            }
            DWORD flags = 0;
            WSARecvFrom(m_sockfd, &(m_ioData.wsabuf), 1, &bytesTransferred, &flags, (SOCKADDR *)&(m_ioData.remoteAddr), &(m_ioData.remoteAddrLen), &(m_ioData.overlapped), NULL);
        }

        // 关闭事件对象
        WSACloseEvent(overlapped.hEvent);
        CloseHandle(m_iocpPort); // 确保关闭 IOCP 句柄

#else
        struct sockaddr_in client_addr;
        socklen_t addr_len = sizeof(client_addr);
#define MAX_EVENTS 10

        int epollFd = epoll_create1(0);
        epoll_event ev;
        ev.events = EPOLLIN | EPOLLET;
        ev.data.fd = m_sockfd;

        epoll_ctl(epollFd, EPOLL_CTL_ADD, m_sockfd, &ev);
        epoll_event events[MAX_EVENTS];

        while (m_isRunning)
        {
            int nfds = epoll_wait(epollFd, events, MAX_EVENTS, -1);
            if (nfds == -1)
            {
                printf("epoll_wait failed\n");
                continue;
            }
            for (int i = 0; i < nfds; ++i)
            {
                if (events[i].data.fd == m_sockfd)
                {
                    int bytes_received = recvfrom(m_sockfd, recvBuffer, sizeof(recvBuffer), 0, (struct sockaddr *)&client_addr, &addr_len);
                    if (bytes_received > 0)
                    {
                        if (onMessage)
                        {
                            onMessage(recvBuffer, bytes_received, (struct sockaddr *)&client_addr);
                        }
                    }
                }
            }
        }

        /* // select
        fd_set readfds;
        int maxfd, activity;
        struct timeval timeout;

        while (m_isRunning)
        {
            FD_ZERO(&readfds);
            FD_SET(m_sockfd, &readfds);
            maxfd = m_sockfd;

            timeout.tv_sec = 5;
            timeout.tv_usec = 0;

            activity = select(maxfd + 1, &readfds, nullptr, nullptr, &timeout);
            if (activity >= 0)
            {
                if (FD_ISSET(m_sockfd, &readfds))
                {
                    int bytes_received = recvfrom(m_sockfd, recvBuffer, sizeof(recvBuffer), 0, (struct sockaddr*)&client_addr, &addr_len);
                    if (bytes_received > 0)
                    {
                        if (onMessage)
                        {
                            onMessage(recvBuffer, bytes_received, (struct sockaddr*)&client_addr);
                        }
                    }
                }
            }
        }
        */

#endif
    }

    void UdpServer::sendData(const void *data, int size, struct sockaddr *client_addr)
    {
#if WIN32
        WSABUF sendBuf{
            .len = (u_long)size,
            .buf = (char *)data};
        WSASendTo(m_sockfd, &sendBuf, 1, nullptr, 0, client_addr, sizeof(struct sockaddr), nullptr, nullptr);
#else
        sendto(m_sockfd, (const char *)data, size, 0, client_addr, sizeof(struct sockaddr));
#endif
    }

    void UdpServer::start()
    {
        if (m_isRunning == false)
        {
            m_isRunning = true;
            m_receiveThread = new std::thread(&UdpServer::receiveData, this);
        }
    }

    void UdpServer::stop()
    {
        if (m_isRunning == true)
        {
            m_isRunning = false;

#ifdef WIN32
            closesocket(m_sockfd);
            CloseHandle(m_iocpPort);
#else
            shutdown(m_sockfd, SHUT_RDWR);
#endif

            m_receiveThread->join();
            delete m_receiveThread;
        }
    }
}
